from __future__ import annotations

import asyncio
import os
import sys
from pathlib import Path, PureWindowsPath
from random import randint
from typing import Any

from aiohttp import ClientConnectorError
from chia_rs import CoinRecord, PrivateKey
from chia_rs.sized_bytes import bytes32
from chia_rs.sized_ints import uint32

from chia.cmds.cmds_util import get_any_service_client
from chia.cmds.start_funcs import async_start
from chia.protocols.outbound_message import NodeType
from chia.server.resolve_peer_info import set_peer_info
from chia.simulator.simulator_full_node_rpc_client import SimulatorFullNodeRpcClient
from chia.util.bech32m import decode_puzzle_hash, encode_puzzle_hash
from chia.util.config import load_config, save_config
from chia.util.errors import KeychainFingerprintExists
from chia.util.keychain import Keychain, bytes_to_mnemonic
from chia.wallet.derive_keys import (
    master_sk_to_farmer_sk,
    master_sk_to_pool_sk,
    master_sk_to_wallet_sk,
    master_sk_to_wallet_sk_unhardened,
)
from chia.wallet.puzzles.p2_delegated_puzzle_or_hidden_puzzle import puzzle_hash_for_pk


def get_ph_from_fingerprint(fingerprint: int, key_id: int = 1) -> bytes32:
    priv_key_and_entropy = Keychain().get_private_key_by_fingerprint(fingerprint)
    if priv_key_and_entropy is None:
        raise Exception("Fingerprint not found")
    private_key = priv_key_and_entropy[0]
    sk_for_wallet_id: PrivateKey = master_sk_to_wallet_sk(private_key, uint32(key_id))
    puzzle_hash: bytes32 = puzzle_hash_for_pk(sk_for_wallet_id.get_g1())
    return puzzle_hash


def create_chia_directory(
    chia_root: Path,
    fingerprint: int,
    farming_address: str | None,
    plot_directory: str | None,
    auto_farm: bool | None,
    docker_mode: bool,
) -> dict[str, Any]:
    """
    This function creates a new chia directory and returns a heavily modified config,
    suitable for use in the simulator.
    """
    from chia.cmds.init_funcs import chia_init

    if not chia_root.is_dir() or not Path(chia_root / "config" / "config.yaml").exists():
        # create chia directories & load config
        chia_init(chia_root, testnet=True, fix_ssl_permissions=True)
        config: dict[str, Any] = load_config(chia_root, "config.yaml")
        # apply standard block-tools config.
        config["full_node"]["send_uncompact_interval"] = 0
        config["full_node"]["target_uncompact_proofs"] = 30
        config["full_node"]["peer_connect_interval"] = 50
        config["full_node"]["sanitize_weight_proof_only"] = False
        config["logging"]["log_level"] = "INFO"  # extra logs for easier development
        # make sure we don't try to connect to other nodes.
        config["full_node"]["introducer_peer"] = None
        config["wallet"]["introducer_peer"] = None
        config["full_node"]["dns_servers"] = []
        config["wallet"]["dns_servers"] = []
        # create custom testnet (simulator0)
        config["network_overrides"]["constants"]["simulator0"] = config["network_overrides"]["constants"][
            "testnet0"
        ].copy()
        config["network_overrides"]["config"]["simulator0"] = config["network_overrides"]["config"]["testnet0"].copy()
        sim_genesis = "eb8c4d20b322be8d9fddbf9412016bdffe9a2901d7edb0e364e94266d0e095f7"
        config["network_overrides"]["constants"]["simulator0"]["GENESIS_CHALLENGE"] = sim_genesis
        config["network_overrides"]["constants"]["simulator0"]["AGG_SIG_ME_ADDITIONAL_DATA"] = sim_genesis
        # tell services to use simulator0
        config["selected_network"] = "simulator0"
        config["wallet"]["selected_network"] = "simulator0"
        config["full_node"]["selected_network"] = "simulator0"
        config["seeder"]["selected_network"] = "simulator0"
        config["harvester"]["selected_network"] = "simulator0"
        config["pool"]["selected_network"] = "simulator0"
        config["farmer"]["selected_network"] = "simulator0"
        config["timelord"]["selected_network"] = "simulator0"
        config["ui"]["selected_network"] = "simulator0"
        config["introducer"]["selected_network"] = "simulator0"
        config["data_layer"]["selected_network"] = "simulator0"
        if not docker_mode:  # We want predictable ports for our docker image.
            # set ports and networks, we don't want to cause a port conflict.
            port_offset = randint(1, 20000)
            config["daemon_port"] -= port_offset
            config["network_overrides"]["config"]["simulator0"]["default_full_node_port"] = 38444 + port_offset
            # wallet
            config["wallet"]["rpc_port"] += port_offset
            # full node
            config["full_node"]["port"] -= port_offset
            config["full_node"]["rpc_port"] += port_offset
            # connect wallet to full node
            set_peer_info(config["wallet"], peer_type=NodeType.FULL_NODE, peer_port=config["full_node"]["port"])
            # ui
            config["ui"]["daemon_port"] = config["daemon_port"]
        else:
            config["self_hostname"] = "0.0.0.0"  # Bind to all interfaces.
            config["logging"]["log_stdout"] = True  # Log to console.
    else:
        config = load_config(chia_root, "config.yaml")
    # simulator overrides
    sim_config = config["simulator"]
    sim_config["key_fingerprint"] = fingerprint
    if farming_address is None:
        prefix = config["network_overrides"]["config"]["simulator0"]["address_prefix"]
        farming_address = encode_puzzle_hash(get_ph_from_fingerprint(fingerprint), prefix)
    sim_config["farming_address"] = farming_address
    if plot_directory is not None:
        sim_config["plot_directory"] = plot_directory
    # Temporary change to fix win / linux differences.
    sim_config["plot_directory"] = str(Path(sim_config["plot_directory"]))
    if "//" in sim_config["plot_directory"] and os.name != "nt":
        # if we're on linux, we need to convert to a linux path.
        sim_config["plot_directory"] = str(PureWindowsPath(sim_config["plot_directory"]).as_posix())
    sim_config["auto_farm"] = auto_farm if auto_farm is not None else True
    farming_ph = decode_puzzle_hash(farming_address)
    # modify genesis block to give the user the reward
    simulator_consts = config["network_overrides"]["constants"]["simulator0"]
    simulator_consts["GENESIS_PRE_FARM_FARMER_PUZZLE_HASH"] = farming_ph.hex()
    simulator_consts["GENESIS_PRE_FARM_POOL_PUZZLE_HASH"] = farming_ph.hex()
    # get fork heights then write back to config
    if "HARD_FORK_HEIGHT" not in sim_config:  # this meh code is done so that we also write to the config file.
        sim_config["HARD_FORK_HEIGHT"] = 0
    simulator_consts["HARD_FORK_HEIGHT"] = sim_config["HARD_FORK_HEIGHT"]

    # save config and return the config
    save_config(chia_root, "config.yaml", config)
    return config


def display_key_info(fingerprint: int, prefix: str) -> None:
    """
    Display key info for a given fingerprint, similar to the output of `chia keys show`.
    """
    print(f"Using fingerprint {fingerprint}")
    private_key_and_seed = Keychain().get_private_key_by_fingerprint(fingerprint)
    if private_key_and_seed is None:
        print(f"Fingerprint {fingerprint} not found")
        return
    sk, seed = private_key_and_seed
    print("\nFingerprint:", sk.get_g1().get_fingerprint())
    print("Master public key (m):", sk.get_g1())
    print("Farmer public key (m/12381/8444/0/0):", master_sk_to_farmer_sk(sk).get_g1())
    print("Pool public key (m/12381/8444/1/0):", master_sk_to_pool_sk(sk).get_g1())
    first_wallet_sk: PrivateKey = master_sk_to_wallet_sk_unhardened(sk, uint32(0))
    wallet_address: str = encode_puzzle_hash(puzzle_hash_for_pk(first_wallet_sk.get_g1()), prefix)
    print(f"First wallet address: {wallet_address}")
    assert seed is not None
    print("Master private key (m):", bytes(sk).hex())
    print("First wallet secret key (m/12381/8444/2/0):", master_sk_to_wallet_sk(sk, uint32(0)))
    mnemonic = bytes_to_mnemonic(seed)
    print("  Mnemonic seed (24 secret words):")
    print(f"{mnemonic} \n")


def generate_and_return_fingerprint(mnemonic: str | None = None) -> int:
    """
    Generate and add new PrivateKey and return its fingerprint.
    """
    from chia.util.keychain import generate_mnemonic

    if mnemonic is None:
        print("Generating private key")
        mnemonic = generate_mnemonic()
    try:
        sk = Keychain().add_key(mnemonic, None)
        fingerprint: int = sk.get_g1().get_fingerprint()
    except KeychainFingerprintExists as e:
        fingerprint = e.fingerprint
        print(f"Fingerprint: {fingerprint} for provided private key already exists.")
        return fingerprint
    print(f"Added private key with public key fingerprint {fingerprint}")
    return fingerprint


def select_fingerprint(
    fingerprint: int | None = None, mnemonic_string: str | None = None, auto_generate_key: bool = False
) -> int | None:
    """
    Either select an existing fingerprint or create one and return it.
    """
    if mnemonic_string:
        fingerprint = generate_and_return_fingerprint(mnemonic_string)
    fingerprints: list[int] = [pk.get_fingerprint() for pk in Keychain().get_all_public_keys()]
    if fingerprint is not None and fingerprint in fingerprints:
        return fingerprint
    elif fingerprint is not None and fingerprint not in fingerprints:
        print(f"Invalid Fingerprint. Fingerprint {fingerprint} was not found.")
        return None
    if auto_generate_key and len(fingerprints) == 1:
        return fingerprints[0]
    if len(fingerprints) == 0:
        if not auto_generate_key:
            if (
                input("No keys in keychain. Press 'q' to quit, or press any other key to generate a new key.").lower()
                == "q"
            ):
                return None
        # generate private key and add to wallet
        fingerprint = generate_and_return_fingerprint()
    else:
        print("Fingerprints:")
        print(
            "If you already used one of these keys, select that fingerprint to skip the plotting process."
            " Otherwise, select any key below."
        )
        for i, fp in enumerate(fingerprints):
            row: str = f"{i + 1}) "
            row += f"{fp}"
            print(row)
        val = None
        prompt: str = f"Choose a simulator key [1-{len(fingerprints)}] ('q' to quit, or 'g' to generate a new key): "
        while val is None:
            val = input(prompt)
            if val == "q":
                return None
            elif val == "g":
                fingerprint = generate_and_return_fingerprint()
                break
            elif not val.isdigit():
                val = None
            else:
                index = int(val) - 1
                if index < 0 or index >= len(fingerprints):
                    print("Invalid value")
                    val = None
                    continue
                else:
                    fingerprint = fingerprints[index]
        assert fingerprint is not None
    return fingerprint


async def generate_plots(config: dict[str, Any], root_path: Path, fingerprint: int, bitfield: bool) -> None:
    """
    Pre-Generate plots for the new simulator instance.
    """

    from chia.simulator.block_tools import BlockTools, test_constants
    from chia.simulator.start_simulator import PLOT_SIZE, PLOTS

    farming_puzzle_hash = decode_puzzle_hash(config["simulator"]["farming_address"])
    os.environ["CHIA_ROOT"] = str(root_path)  # change env variable, to make it match what the daemon would set it to

    # create block tools and use local keychain
    bt = BlockTools(
        test_constants,
        root_path,
        automated_testing=False,
        plot_dir=config["simulator"].get("plot_directory", "plots"),
        keychain=Keychain(),
    )
    await bt.setup_keys(fingerprint=fingerprint, reward_ph=farming_puzzle_hash)
    existing_plots = await bt.setup_plots(
        num_og_plots=PLOTS, num_pool_plots=0, num_non_keychain_plots=0, plot_size=PLOT_SIZE, bitfield=bitfield
    )
    print(f"{'New plots generated.' if existing_plots else 'Using Existing Plots'}\n")


async def get_current_height(root_path: Path) -> int:
    async with get_any_service_client(SimulatorFullNodeRpcClient, root_path=root_path, consume_errors=False) as (
        node_client,
        _,
    ):
        num_blocks = len(await node_client.get_all_blocks())
    return num_blocks


async def async_config_wizard(
    root_path: Path,
    fingerprint: int | None,
    farming_address: str | None,
    plot_directory: str | None,
    mnemonic_string: str | None,
    auto_farm: bool | None,
    docker_mode: bool,
    bitfield: bool,
) -> None:
    # either return passed through fingerprint or get one
    fingerprint = select_fingerprint(fingerprint, mnemonic_string, docker_mode)
    if fingerprint is None:
        # user cancelled wizard
        return
    # create chia directory & get config.
    print("Creating chia directory & config...")
    config = create_chia_directory(root_path, fingerprint, farming_address, plot_directory, auto_farm, docker_mode)
    # Pre-generate plots by running block_tools init functions.
    print("Please Wait, Generating plots...")
    print("This may take up to a minute if you are on a slow machine")

    await generate_plots(config, root_path, fingerprint, bitfield)
    # final messages
    final_farming_address = config["simulator"]["farming_address"]
    print(f"\nFarming & Prefarm reward address: {final_farming_address}\n")
    print("Configuration Wizard Complete.")
    print("Starting Simulator now...\n\n")

    sys.argv[0] = str(Path(sys.executable).parent / "chia")  # fix path for tests
    await async_start(root_path, config, ("simulator",), restart=False, skip_keyring=False)

    # now we make sure the simulator has a genesis block
    print("Please wait, generating genesis block.")
    while True:
        try:
            num_blocks: int = await get_current_height(root_path)
        except ClientConnectorError:
            await asyncio.sleep(0.25)
        else:
            if num_blocks == 0:
                await farm_blocks(None, root_path, 1, True, final_farming_address)
                print("Genesis block generated, exiting.")
            else:
                print("Genesis block already exists, exiting.")
            break
    print(f"\nMake sure your CHIA_ROOT Environment Variable is set to: {root_path}")


def print_coin_record(
    name: str,
    address_prefix: str,
    coin_record: CoinRecord,
) -> None:
    from datetime import datetime

    coin_address = encode_puzzle_hash(coin_record.coin.puzzle_hash, address_prefix)
    print(f"Coin 0x{coin_record.name.hex()}")
    print(f"Wallet Address: {coin_address}")
    print(f"Confirmed at block: {coin_record.confirmed_block_index}")
    print(f"Spent: {f'at Block {coin_record.spent_block_index}' if coin_record.spent else 'No'}")
    print(f"Coin Amount: {coin_record.coin.amount} {name}")
    print(f"Parent Coin ID: 0x{coin_record.coin.parent_coin_info.hex()}")
    print(f"Created at: {datetime.fromtimestamp(float(coin_record.timestamp)).strftime('%Y-%m-%d %H:%M:%S')}\n")


async def print_coin_records(
    config: dict[str, Any],
    node_client: SimulatorFullNodeRpcClient,
    include_reward_coins: bool,
    include_spent: bool = False,
) -> None:
    import sys

    coin_records: list[CoinRecord] = await node_client.get_all_coins(include_spent)
    coin_records = [coin_record for coin_record in coin_records if not coin_record.coinbase or include_reward_coins]
    address_prefix = config["network_overrides"]["config"][config["selected_network"]]["address_prefix"]
    name = "mojo"
    paginate = False  # I might change this later.
    if len(coin_records) != 0:
        print("All Coins: ")
        if paginate is True:
            paginate = sys.stdout.isatty()
        num_per_screen = 5 if paginate else len(coin_records)
        # ripped from cmds/wallet_funcs.
        for i in range(0, len(coin_records), num_per_screen):
            for j in range(num_per_screen):
                if i + j >= len(coin_records):
                    break
                print_coin_record(
                    coin_record=coin_records[i + j],
                    name=name,
                    address_prefix=address_prefix,
                )
            if i + num_per_screen <= len(coin_records) and paginate:
                print("Press q to quit, or c to continue")
                while True:
                    entered_key = sys.stdin.read(1)
                    if entered_key == "q":
                        return None
                    elif entered_key == "c":
                        break


async def print_wallets(config: dict[str, Any], node_client: SimulatorFullNodeRpcClient) -> None:
    ph_and_amount = await node_client.get_all_puzzle_hashes()
    address_prefix = config["network_overrides"]["config"][config["selected_network"]]["address_prefix"]
    name = "mojo"
    for puzzle_hash, (amount, num_tx) in ph_and_amount.items():
        address = encode_puzzle_hash(puzzle_hash, address_prefix)
        print(f"Address: {address} has a balance of: {amount} {name}, with a total of: {num_tx} transactions.\n")


async def print_status(
    rpc_port: int | None,
    root_path: Path,
    fingerprint: int | None,
    show_key: bool,
    show_coins: bool,
    include_reward_coins: bool,
    show_addresses: bool,
) -> None:
    """
    This command allows users to easily get the status of the simulator
    and information about the state of and the coins in the simulated blockchain.
    """
    from chia.cmds.show_funcs import print_blockchain_state
    from chia.cmds.units import units

    async with get_any_service_client(SimulatorFullNodeRpcClient, root_path, rpc_port) as (node_client, config):
        # Display keychain info
        if show_key:
            if fingerprint is None:
                fingerprint = config["simulator"]["key_fingerprint"]
            if fingerprint is not None:
                display_key_info(
                    fingerprint, config["network_overrides"]["config"][config["selected_network"]]["address_prefix"]
                )
            else:
                print(
                    "No fingerprint in config, either rerun 'cdv sim create' "
                    "or use --fingerprint to specify one, skipping key information."
                )
        # chain status ( basically chia show -s)
        await print_blockchain_state(node_client, config)
        print("")
        # farming information
        target_ph: bytes32 = await node_client.get_farming_ph()
        farming_coin_records = await node_client.get_coin_records_by_puzzle_hash(target_ph, False)
        prefix = config["network_overrides"]["config"][config["selected_network"]]["address_prefix"]
        print(
            f"Current Farming address: {encode_puzzle_hash(target_ph, prefix)}, "
            f"with a balance of: "
            f"{sum(coin_records.coin.amount for coin_records in farming_coin_records) / units['chia']} TXCH."
        )
        if show_addresses:
            print("All Addresses: ")
            await print_wallets(config, node_client)
        if show_coins:
            await print_coin_records(config, node_client, include_reward_coins)


async def revert_block_height(
    rpc_port: int | None,
    root_path: Path,
    num_blocks: int,
    num_new_blocks: int,
    reset_chain_to_genesis: bool,
    use_revert_blocks: bool,
) -> None:
    """
    This function allows users to easily revert the chain to a previous state or perform a reorg.
    """
    async with get_any_service_client(SimulatorFullNodeRpcClient, root_path, rpc_port) as (node_client, _):
        if use_revert_blocks:
            if num_new_blocks != 1:
                print(f"Ignoring num_new_blocks: {num_new_blocks}, because we are not performing a reorg.")
            # in this case num_blocks is the number of blocks to delete
            new_height: int = await node_client.revert_blocks(num_blocks, reset_chain_to_genesis)
            print(
                f"All transactions in Block: {new_height + num_blocks} and above were successfully deleted, "
                "you should now delete & restart all wallets."
            )
        else:
            # However, in this case num_blocks is the fork height.
            new_height = await node_client.reorg_blocks(num_blocks, num_new_blocks, use_revert_blocks)
            old_height = new_height - num_new_blocks
            print(f"All transactions in Block: {old_height - num_blocks} and above were successfully reverted.")
        print(f"Block Height is now: {new_height}")


async def farm_blocks(
    rpc_port: int | None,
    root_path: Path,
    num_blocks: int,
    transaction_blocks: bool,
    target_address: str,
) -> None:
    """
    This function is used to generate new blocks.
    """
    async with get_any_service_client(SimulatorFullNodeRpcClient, root_path, rpc_port) as (node_client, config):
        if target_address == "":
            target_address = config["simulator"]["farming_address"]
        if target_address is None:
            print(
                "No target address in config, falling back to the temporary address currently in use. "
                "You can use 'cdv sim create' or use --target-address to specify a different address."
            )
            target_ph: bytes32 = await node_client.get_farming_ph()
        else:
            target_ph = decode_puzzle_hash(target_address)
        await node_client.farm_block(target_ph, num_blocks, transaction_blocks)
        print(f"Farmed {num_blocks}{' Transaction' if transaction_blocks else ''} blocks")
        block_height = (await node_client.get_blockchain_state())["peak"].height
        print(f"Block Height is now: {block_height}")


async def set_auto_farm(rpc_port: int | None, root_path: Path, set_autofarm: bool) -> None:
    """
    This function can be used to enable or disable Auto Farming.
    """
    async with get_any_service_client(SimulatorFullNodeRpcClient, root_path, rpc_port) as (node_client, _):
        current = await node_client.get_auto_farming()
        if current == set_autofarm:
            print(f"Auto farming is already {'on' if set_autofarm else 'off'}")
            return
        result = await node_client.set_auto_farming(set_autofarm)
        print(f"Auto farming is now {'on' if result else 'off'}")
